import os
import unittest

import torch
from torch.backends import cudnn
from torchsummary import summary

from models.backbone.xception import aligned_xception

torch.backends.cudnn.benchmark = True

os.environ["CUDA_VISIBLE_DEVICES"] = '1'


class XceptionTestCase(unittest.TestCase):

    def test_summary_xception(self):
        self.assertTrue(torch.cuda.is_available())
        model = aligned_xception(pretrained=True)

        inputs = torch.rand(1, 3, 512, 512)
        output = model(inputs)
        print(output[0].size(), output[1].size())

        model = model.cuda()
        input_size = (3, 512, 512)
        summary(model, input_size=input_size)


if __name__ == '__main__':
    unittest.main()
